#include <stdio.h> 
#include <omp.h>
#include <Windows.h>

double* alloc_init_B(double* A, int N) {
    double* B = new double[N];
    for (int i = 0; i < N; i++) {
        B[i] = A[i];
    }

    // 输出线程 ID 和所在 CPU 核心
    printf("[alloc_init_B] Thread ID: %lu, CPU core: %lu\n",
           GetCurrentThreadId(), GetCurrentProcessorNumber());

    return B;
}

void compute_on_B(double* B, int N) {
    for (int i = 0; i < N; i++) {
        B[i] = B[i] * 2;
    }

    // 输出线程 ID 和所在 CPU 核心
    printf("[compute_on_B] Thread ID: %lu, CPU core: %lu\n",
           GetCurrentThreadId(), GetCurrentProcessorNumber());
}

void task_affinity(double* A, int N) {
    double* B;

    #pragma omp parallel
    #pragma omp single
    {
        #pragma omp task depend(out:B) shared(B) affinity(A)
        {
            B = alloc_init_B(A, N);
        }

        #pragma omp task depend(in:B) shared(B) affinity(A)
        {
            compute_on_B(B, N);
        }

        #pragma omp taskwait
    }

    // 打印结果，查看 B 数组的内容
    for (int i = 0; i < N; i++) {
        printf("B[%d] = %f\n", i, B[i]);
    }
}

int main() {
    int N = 5;
    double A[] = {1, 2, 3, 4, 5};

    task_affinity(A, N);

    return 0;
}
